# -*- coding: utf-8 -*-
"""
网络设备自动化巡检脚本（Ollama 本地大模型版）
功能说明：
1. 通过 SSH 批量执行网络设备巡检命令
2. 使用本地部署的 Ollama 大模型分析日志
3. 生成带时间戳的完整分析报告
"""

# ​********************​ 基础库导入 ​********************
import paramiko  # SSH 连接库[6](@ref)
import os
import time
import sys
import socket  # 服务检测
from datetime import datetime
import ollama  # Ollama 本地大模型客户端库[1,4](@ref)

# ​********************​ 配置区域 ​********************
SSH_PORT = 22  # SSH 端口号（默认值）
SSH_USER = "sysadmin"  # 设备登录用户名
SSH_PASS = "Admin@h3c"  # 设备登录密码
OLLAMA_MODEL = "deepseek-r1:14b"  # 本地部署的大模型名称[4](@ref)
OLLAMA_ENDPOINT = "http://localhost:11434"  # Ollama 服务地址[5](@ref)
COMMAND_DELAY = 2  # 命令执行间隔（防设备过载）
TIMEOUT = 30  # SSH 连接超时时间
INSPECTION_COMMANDS = [
    'display version',
    'display cpu-usage'
]
# ​*************************************************

def get_timestamp():
    """生成标准化时间戳（用于日志记录）"""
    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")

def collect_device_info(ip):
    """
    SSH 设备信息采集函数
    :param ip: 目标设备IP地址
    :return: 结构化采集结果（含错误处理）
    """
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())  # 自动接受密钥
    
    try:
        # 建立 SSH 连接
        print(f"[{get_timestamp()}] 正在连接 {ip}...")
        client.connect(ip, SSH_PORT, SSH_USER, SSH_PASS, timeout=TIMEOUT)
        output = f"\n=== {ip} 巡检结果 ===\n"

        # 执行预定义巡检命令
        for idx, cmd in enumerate(INSPECTION_COMMANDS):
            try:
                print(f"[{get_timestamp()}] 执行命令: {cmd}")
                stdin, stdout, stderr = client.exec_command(cmd)
                result = stdout.read().decode('utf-8', errors='ignore').strip()
                output += f"[{cmd}]\n{result}\n\n"

                # 非最后命令添加保护间隔
                if idx < len(INSPECTION_COMMANDS) - 1:
                    print(f"[{get_timestamp()}] 等待 {COMMAND_DELAY} 秒...")
                    time.sleep(COMMAND_DELAY)

            except Exception as cmd_error:
                output += f"命令执行失败: {str(cmd_error)}\n"

        return output

    except Exception as conn_error:
        return f"\n=== {ip} 连接失败 ===\n错误信息: {str(conn_error)}\n"
    finally:
        client.close()  # 确保连接关闭

def generate_ai_report(content):
    """
    Ollama 大模型报告生成函数
    :param content: 原始巡检数据
    :return: 结构化分析报告（含错误处理）
    """
    try:
        print(f"[{get_timestamp()}] 正在生成分析报告...")
        response = ollama.chat(
            model=OLLAMA_MODEL,
            messages=[
                {"role": "system", "content": "你是一个专业的网络设备分析专家，能够根据设备巡检日志生成分析报告"},
                {"role": "user", "content": f"请根据以下设备巡检日志生成分析报告：\n{content}"}
            ]
        )
        return response['message']['content'].strip()
    except Exception as e:
        return f"Ollama 服务异常：{str(e)}"

def check_ollama_service():
    """
    Ollama 服务状态检测函数
    确保本地 Ollama 服务已正常启动[5](@ref)
    """
    try:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.settimeout(3)
            s.connect(('localhost', 11434))
    except Exception as e:
        print(f"[{get_timestamp()}] 错误：Ollama 服务未启动！请执行：ollama run {OLLAMA_MODEL}")
        sys.exit(1)

def main():
    """主执行流程"""
    # 生成带时间戳的报告文件
    report_file = f"Analysis_Report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md"
    
    try:
        # 读取设备IP列表
        with open('ip.txt', 'r') as f:
            ips = [line.strip() for line in f if line.strip()]
        
        if not ips:
            print(f"[{get_timestamp()}] 错误：ip.txt中没有找到有效IP地址")
            return

        # 执行设备巡检
        full_report = ""
        for ip in ips:
            print(f"\n[{get_timestamp()}] 开始巡检 {ip}")
            full_report += collect_device_info(ip)

        # 生成中间文件
        with open('temp.txt', 'w') as f:
            f.write(full_report)

        # 调用大模型分析
        with open('temp.txt', 'r') as f:
            ai_report = generate_ai_report(f.read())

        # 保存最终报告
        with open(report_file, 'w') as f:
            f.write(f"=== 设备巡检分析报告 ===\n")
            f.write(f"生成时间: {get_timestamp()}\n")
            f.write(f"分析模型: {OLLAMA_MODEL}\n")
            f.write(f"巡检设备: {len(ips)}台\n\n")
            f.write(ai_report)

        print(f"\n[{get_timestamp()}] 报告生成成功: {report_file}")

    except FileNotFoundError:
        print(f"[{get_timestamp()}] 错误：未找到ip.txt文件")
    except Exception as e:
        print(f"[{get_timestamp()}] 发生未预期错误: {str(e)}")
    finally:
        # 清理临时文件
        time.sleep(10)
        if os.path.exists('temp.txt'):
            os.remove('temp.txt')
            print(f"[{get_timestamp()}] 已清理临时文件")

if __name__ == "__main__":
    check_ollama_service()  # 服务预检
    main()